import os
import sys

import onnx
import torch
from onnxsim import simplify

sys.path.append('.')
from models.model_stages import BiSeNet

if __name__ == "__main__":
    respth = r'checkpoints/train_STDC1-Seg/train_0418_sigmoid/pths/model_final.pth'

    num_classes = 1
    backbone='STDCNet813'
    inputsize=[800, 800]
    use_boundary_2=False
    use_boundary_4=False
    use_boundary_8=False
    use_boundary_16=False
    use_conv_last=False


    net = BiSeNet(backbone=backbone, n_classes=num_classes,
                use_boundary_2=use_boundary_2, use_boundary_4=use_boundary_4, 
                use_boundary_8=use_boundary_8, use_boundary_16=use_boundary_16, 
                use_conv_last=use_conv_last,
                export=True)
    
    net.load_state_dict(torch.load(respth))
    net.eval()


    dummy_input = torch.randn((1, 3, 800, 800))
    output = net(dummy_input)
    output_path = "checkpoints/export_model/stdc.onnx"

    torch.onnx.export(
        net,
        dummy_input,
        output_path,
        export_params=True, 
        # do_constant_folding=True,
        input_names=["data"],
        output_names=["output"],
        dynamic_axes=None,
        opset_version=11,
        verbose=True
    )

    onnx_model = onnx.load(output_path)
    onnx.checker.check_model(onnx_model)
    model_simp, check = simplify(onnx_model,
                                     dynamic_input_shape=False,
                                     input_shapes=None)
    assert check, "Simplified ONNX model could not be validated"

    output_name = "checkpoints/export_model/stdc_sim.onnx"
    onnx.save(model_simp, output_name)

    
    print("export success !!!")
